#!/usr/bin/python
#coding:utf-8

# Copyright 2012 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, extract the edgels, and find
## the orientation using SQP.
##

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from numpy import dot

from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

from quaternion import Quat, random_quaternion

import simplejson

import Image

set_printoptions(precision=3)

from pylab import *

from scipy.misc import factorial as fact
from sqp_step1 import create_pic


def circle_segment_area(x1,x2,y1,y2):    
    if x1>x2:
        x1,x2=x2,x1
    if y1>y2:
        y1,y2=y2,y1

    xm = (x1+x2)/2
    ym = (y1+y2)/2

    if xm < 0:
        x1,x2 = -x2,-x1
        xm = (x1+x2)/2
    if ym < 0:
        y1,y2 = -y2,-y1
        ym = (y1+y2)/2

    if xm > ym:
        x1,x2,y1,y2 = y1,y2,x1,x2
        xm = (x1+x2)/2
        ym = (y1+y2)/2

    ## Handle the easy cases
    if (x2*x2+y2*y2) <= 1:
        return (x2-x1)*(y2-y1)
    elif (x1*x1+y1*y1) >= 1:
        return 0

    ## Now go on to calculate the intermediary ones...

    ## Caso chato
    if x1<0:
        yc2 = sqrt((1-x2*x2))
        xc1 = x2 * y1/yc2

        ## Circular sector area
        # (0,1) (x2,yc2)
        area = arcsin(x2)
        ## Now subtract the outer triangles
        area -= xc1*y1
        ## Add inner triangles
        area += (x2-xc1)*(yc2-y1)
        return area

    if x1*x1+y2*y2 >= 1 and x2*x2+y1*y1 >= 1 :
        yc1 = sqrt((1-x1*x1))
        xc1 = sqrt((1-y1*y1))

        ## Circular sector area
        ## (x1,yc1) (xc1,y1)
        area = arcsin(xc1*yc1 - x1*y1)
        ## Now subtract the outer triangles
        area -= (xc1-x1)*y1 + (yc1-y1)*x1
        area = area/2
        return area


    if (x1*x1+y2*y2) >= 1 and (x2*x2+y1*y1) <= 1 :
        yc1 = sqrt((1-x1*x1))
        yc2 = sqrt((1-x2*x2))
        xc1 = x2 * y1/yc2

        ## Circular sector area
        # (x1,yc1) (x2,yc2)
        area = arcsin(x2*yc1 - x1*yc2)
        ## Now subtract the outer triangles
        area -= (xc1-x1)*y1 + (yc1-y1)*x1
        ## Add inner triangles
        area += (x2-xc1)*(yc2-y1)
        area = area/2
        return area


    #else:
    if x1*x1+y2*y2 <= 1 and x2*x2+y1*y1 <= 1 :
        yc2 = sqrt((1-x2*x2))
        xc2 = sqrt((1-y2*y2))
        yc1 = y2 * x1/xc2
        xc1 = x2 * y1/yc2

        ## Circular sector area
        # (xc2, y2) (x2, yc2)
        area = arcsin(x2*y2 - xc2*yc2)
        ## Now subtract the outer triangles
        area -= (xc1-x1)*y1 + (yc1-y1)*x1
        ## Add inner triangles
        area += (x2-xc1)*(yc2-y1) + (y2-yc1)*(xc2-x1)
        area = area/2
        return area

    raise Exception('Irregular case')




def circular_mask(N):
    out = zeros((N,N))
    
    for j in mgrid[0.:N]:
        for k in mgrid[0.:N]:
            x1 = (k*2.0-N) / N
            x2 = ((k+1)*2.0-N) / N
            y1 = (j*2.0-N) / N
            y2 = ((j+1)*2.0-N) / N

            out[j,k] = circle_segment_area(x1,x2,y1,y2)
    return out


# Formulas from "A Moment-Based Unified Approach to Image Feature
# Detection" by Sugata Ghosal and Rajiv Mehrotra
def zernike_Rnm(n,m,rho):
    m = np.abs(m)
    assert n>=m
    assert (n-m)%2==0

    out = zeros(rho.shape)
    for s in range(1+(n-m)/2):
        out += ( (-1.0)**s * fact(n-s)*rho**(n-2*s) /
                 (fact(s) * fact((n+m)/2-s) * fact((n-m)/2-s)))
    return out

def zernike_mask(N,n,m):
    out = 0j * ones((N,N))
    for j in range(N):
        for k in range(N):
            
            x = (k*2.0+1-N) / N
            y = (j*2.0+1-N) / N
            z = x+1j*y
            R = zernike_Rnm(n,m,np.abs(z))
            out[j,k] = (n+1)/pi * R * exp(1j * m * np.angle(z))
    cm = circular_mask(N)
    return cm * out


def plota_polis():
    ## Display the radial polinomials
    x = mgrid[0:1.0:.001]
    figure(1)
    title('Radial polinomials')
    for n,m in pars:
        y = array([zernike_Rnm(n,m,xx) for xx in x])
        plot(x,y, label='-')
    axis([-0.05,1.05,-1.05,1.05])
    grid()

def plota_masks():
    figure(2)
    title('Masks')
    vl=.07
    for n in range(4):
        Vnm = VV[n]
        subplot(3,4,1+2*n)
        imshow(real(Vnm), vmin=-vl, vmax=vl)
        subplot(3,4,2+2*n)
        imshow(imag(Vnm), vmin=-vl, vmax=vl)


def heavi(x,l):
    # return clip(x*2,-1,1)
    return clip(x*.9 + l*3.5,-1,1)

def render_step(Nw, ang,h1,h2,l):
    mystep = zeros((Nw,Nw))

    dx = cos(ang)
    dy = sin(ang)

    for j in range(Nw):
        for k in range(Nw):
            x,y = (k-3.0), (j-3.0)
            c = heavi(x*dx+y*dy, l)
            mystep[j,k] = c*h1 + (1-c)*h2
    return mystep











def basic_test_synthesize_img():    
    ion()

    pars = [(0,0),(1,1),(2,0),(3,1)]

    Nw = 7

    VV=[]

    yingdong_V11 = array([[ 00,  15,  19, 0, -19, -15, -00],
                          [ 22,  47,  23, 0, -23, -47, -22],
                          [ 57,  47,  23, 0, -23, -47, -57],
                          [ 70,  47,  23, 0, -23, -47, -70],
                          [ 57,  47,  23, 0, -23, -47, -57],
                          [ 22,  47,  23, 0, -23, -47, -22],
                          [ 00,  15,  19, 0, -19, -19, -00]])*1e-3

    yingdong_V20 = array([[ 00,  23,  39,  41,  39,  23,  00],
                          [ 23,  27, -13, -26, -13,  27,  23],
                          [ 39, -13, -53, -66, -53, -13,  39],
                          [ 41, -26, -66, -81, -66, -26,  41],
                          [ 39, -13, -53, -66, -53, -13,  39],
                          [ 23,  27, -13, -26, -13,  27,  23],
                          [ 00,  23,  39,  41,  39,  23,  00]])*1e-3


    Vnm = zernike_mask(Nw,0,0)
    VV.append(real(Vnm))
    #Vnm = zernike_mask(Nw,1,1)
    #VV.append(-(Vnm))
    VV.append(yingdong_V11 +1j*yingdong_V11.T)
    # Vnm = zernike_mask(Nw,2,0)
    # VV.append(real(Vnm))

    VV.append(yingdong_V20)

    Vnm = zernike_mask(Nw,3,1)
    VV.append(Vnm)



    ol = .3
    oh1 = 1.0
    oh2 = -1.0

    Ntry = 10000

    moL = 0j * zeros((Ntry, 4))
    parL = zeros((Ntry, 3))


    figure(10)
    for N in range(Ntry):
        oang = rand()*2*pi-pi
        ol = rand()*sqrt(2.2)/7

        parL[N] = [oang, ol, oh2]

        mystep = render_step(Nw, oang,oh1,oh2,ol)
        mystep += (random(mystep.shape)-.5) * 2 * .2

        if N < 9:
            subplot(3,3,N+1)
            imshow(mystep)


        ## Calculate moments
        AA = []
        for n in range(4):
            Vnm = VV[n]
            Anm = ((Vnm) * mystep).sum()
            AA.append(Anm)

        moL[N] = AA

    ## recovered params
    hatL = zeros((Ntry,3))
    for N in range(Ntry):
        A00,A11,A20,A31 = moL[N]
        A00 = np.abs(A00)
        A20 = np.abs(A20)
        # print 70*'--'
        # print A00,A11,A20,A31
        # print parL[N]

        A31_ = np.abs(A31)

        th = (pi+np.angle(-(A11)))%(2*pi)-pi

        # l = ((sqrt(25. * A31_**2 + 24. * A20**2) + 5.*A31_ ) /
        #      (12. * A20))
        l = A20/np.abs(A11)

        h = -15 * A31_ / ( 2 * (1-l)**1.5 * (1-6*l**2)  )

        # print th, l, h

        hatL[N] = [th,l,h]


    figure(1)
    ii = argsort(hatL[:,0])
    plot(parL[ii,0]*180/pi, (parL[ii,0]-hatL[ii,0])*180/pi, ',')
    grid()

    figure(2)
    ii = argsort(hatL[:,1])
    plot(3.5 * parL[ii,1], 3.5 * (parL[ii,1]-hatL[ii,1]), ',')
    grid()





if __name__ == '__main__':



    V11 = array([[ 00,  15,  19, 0, -19, -15, -00],
                 [ 22,  47,  23, 0, -23, -47, -22],
                 [ 57,  47,  23, 0, -23, -47, -57],
                 [ 70,  47,  23, 0, -23, -47, -70],
                 [ 57,  47,  23, 0, -23, -47, -57],
                 [ 22,  47,  23, 0, -23, -47, -22],
                 [ 00,  15,  19, 0, -19, -19, -00]])*1e-3

    V20 = array([[ 00,  23,  39,  41,  39,  23,  00],
                 [ 23,  27, -13, -26, -13,  27,  23],
                 [ 39, -13, -53, -66, -53, -13,  39],
                 [ 41, -26, -66, -81, -66, -26,  41],
                 [ 39, -13, -53, -66, -53, -13,  39],
                 [ 23,  27, -13, -26, -13,  27,  23],
                 [ 00,  23,  39,  41,  39,  23,  00]])*1e-3

    shigeru_filter=scipy.array([
      [ -0.003776, -0.010199, 0., 0.010199, 0.003776 ],
      [ -0.026786, -0.070844, 0., 0.070844, 0.026786 ],
      [ -0.046548, -0.122572, 0., 0.122572, 0.046548 ],
      [ -0.026786, -0.070844, 0., 0.070844, 0.026786 ],
      [ -0.003776, -0.010199, 0., 0.010199, 0.003776 ]
      ])


    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load images and initialize pic objects

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])

    pic = create_pic(job_params, framenum)

    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold'] # ** -1
    ## Run the grid sweep, and generate the edgels list.
    ## Edgel extractor parameters
    pic.extract_edgels(gspc, glim)




    #19 152
    #326 57


    frm = pic.frame[:,:,1]

    imshow(frm/256., cmap=cm.bone)

    

    pp = []
    vv = array([-1,1])
    for ex in range(19, 327, 1):
        ey = 152 + (ex - 19) * (57 - 152) / (326.-19)

        snp = frm[ey-3:ey+4,ex-3:ex+4]

        A11r = (V11 * snp).sum()
        A11i = (-V11.T * snp).sum()
        A20 = (V20 * snp).sum()

        gx = (shigeru_filter * snp[1:-1,1:-1]).sum()
        gy = (shigeru_filter.T * snp[1:-1,1:-1]).sum()

        print 70*'-'
        print A11r, A11i, A20
        
        #th = (pi+np.angle(-(A11)))%(2*pi)-pi
        th = arctan2(A11i, A11r)
        th2 = arctan2(-gy, gx)

        # l = ((sqrt(25. * A31_**2 + 24. * A20**2) + 5.*A31_ ) /
        #      (12. * A20))
        amp = sqrt(A11r+A11r + A11i*A11i)
        l = -3.5 * A20/amp
        
        print th, l

        dx = A11r/amp
        dy = -A11i/amp

        cx = ex
        cy = ey
        plot(ex,ey, 'b+')
        plot(cx+7*dy*vv, cy-7*dx*vv, 'b-')

        cx = ex + l * dx
        cy = ey + l * dy
        plot(cx+7*dy*vv, cy-7*dx*vv, 'r-')

        pp.append((th,th2, l ))


    pp = array(pp)

    Np = pp.shape[0]

    m0 = sort(pp[:,0])[Np/2]
    m1 = sort(pp[:,1])[Np/2]

    xx = mgrid[0.:Np]/Np

    figure(2)
    plot(87*np.sort(np.abs(pp[:,0]-m0)),xx)
    plot(87*np.sort(np.abs(pp[:,1]-m1)),xx)
    grid()

    figure(3)
    plot(pp[:,2],'-+')
    ylim(1,-1)
    grid()

